import React, { useState, useEffect, useRef } from 'react';
import {
    Layout, Card, Button, Input, Form, Space, Typography,
    Divider, Badge, Tag, Alert, Tabs, Row, Col, Spin, message,
    Modal, Collapse, Select
} from 'antd';
import './ChatBox.css';

const { Header, Sider, Content } = Layout;
const { Title, Text } = Typography;
const { TextArea } = Input;
const { Panel } = Collapse;

const ChatBox = () => {
    const [connectionStatus, setConnectionStatus] = useState(false);
    const [serverConfig, setServerConfig] = useState({
        serverName: '',
        serverUri: 'ws://localhost:9000'
    });
    const [connectedServers, setConnectedServers] = useState([]);
    const [currentServer, setCurrentServer] = useState(null);
    const [tools, setTools] = useState([]);
    const [resources, setResources] = useState([]);
    const [messages, setMessages] = useState([]);
    const [loading, setLoading] = useState(false);
    const [inputMessage, setInputMessage] = useState('');
    const [selectedTool, setSelectedTool] = useState(null);
    const [toolParams, setToolParams] = useState('{}');
    const messagesEndRef = useRef(null);
    const [websocket, setWebsocket] = useState(null);
    const [sessionId, setSessionId] = useState(null);

    useEffect(() => {
        scrollToBottom();
    }, [messages]);

    const scrollToBottom = () => {
        messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
    };

    const addMessage = (type, content) => {
        const newMessage = {
            id: Date.now() + Math.random(), // 确保唯一性
            type,
            content,
            timestamp: new Date().toLocaleTimeString()
        };
        setMessages(prev => [...prev, newMessage]);
    };

    const handleConnect = async () => {
        setLoading(true);
        try {
            // 创建WebSocket连接到MCP客户端
            const ws = new WebSocket('ws://localhost:8080/ws/mcp');
            
            ws.onopen = () => {
                setWebsocket(ws);
                // 生成客户端sessionId
                const clientSessionId = 'client-' + Date.now() + '-' + Math.random().toString(36).substr(2, 9);
                setSessionId(clientSessionId);
                // 发送连接请求到MCP服务器
                ws.send(JSON.stringify({
                    type: 'connect',
                    sessionId: clientSessionId,
                    serverUri: serverConfig.serverUri
                }));
                // 存储clientSessionId以便在消息处理中使用
                ws.clientSessionId = clientSessionId;
            };
            
            ws.onmessage = (event) => {
                const response = JSON.parse(event.data);
                handleWebSocketMessage(response);
            };
            
            ws.onerror = (error) => {
                addMessage('error', `WebSocket连接错误: ${error.message}`);
                message.error('连接错误');
                setLoading(false);
            };
            
            ws.onclose = () => {
                setConnectionStatus(false);
                setWebsocket(null);
                addMessage('system', 'WebSocket连接已关闭');
            };
            
        } catch (error) {
            addMessage('error', `连接错误: ${error.message}`);
            message.error('连接错误');
            setLoading(false);
        }
    };

    const handleWebSocketMessage = (response) => {
        switch (response.type) {
            case 'connect':
                if (response.status === 'success') {
                    setConnectionStatus(true);
                    // Use sessionId from the response or from the current sessionId state
                    const clientSessionId = response.sessionId || sessionId;
                    console.log('Connection success - websocket:', websocket);
                    console.log('Connection success - clientSessionId:', clientSessionId);
                    console.log('Connection success - response.sessionId:', response.sessionId);
                    console.log('Connection success - sessionId state:', sessionId);
                    
                    const serverInfo = {
                        uri: serverConfig.serverUri,
                        name: serverConfig.serverName || `服务器 ${connectedServers.length + 1}`,
                        sessionId: clientSessionId,
                        connectedAt: new Date().toLocaleTimeString(),
                        tools: [],
                        resources: []
                    };
                    console.log('Created serverInfo:', serverInfo);
                    
                    setConnectedServers(prev => [...prev, serverInfo]);
                    setCurrentServer(serverInfo);
                    console.log('Set currentServer to:', serverInfo);
                    
                    addMessage('system', `连接MCP服务器成功: ${serverConfig.serverName || serverConfig.serverUri}`);
                    message.success('连接成功');
                    // 连接成功后自动获取工具和资源列表，传递sessionId
                    console.log('About to call loadTools and loadResources with clientSessionId:', clientSessionId);
                    loadTools(clientSessionId);
                    loadResources(clientSessionId);
                } else {
                    addMessage('error', `连接失败: ${response.message}`);
                    message.error('连接失败');
                }
                setLoading(false);
                break;
            case 'disconnect':
                setConnectionStatus(false);
                if (currentServer) {
                    setConnectedServers(prev => prev.filter(server => server.uri !== currentServer.uri));
                }
                setCurrentServer(null);
                setTools([]);
                setResources([]);
                addMessage('system', '已断开连接');
                message.success('断开连接成功');
                setLoading(false);
                break;
            case 'list_tools':
                if (response.status === 'success') {
                    const toolsData = response.data || [];
                    setTools(toolsData);
                    // 更新当前服务器的工具列表
                    if (currentServer) {
                        setConnectedServers(prev => prev.map(server => 
                            server.uri === currentServer.uri 
                                ? { ...server, tools: toolsData }
                                : server
                        ));
                        setCurrentServer(prev => ({ ...prev, tools: toolsData }));
                    }
                    addMessage('system', `发现 ${toolsData.length} 个工具`);
                }
                break;
            case 'list_resources':
                if (response.status === 'success') {
                    const resourcesData = response.data || [];
                    setResources(resourcesData);
                    // 更新当前服务器的资源列表
                    if (currentServer) {
                        setConnectedServers(prev => prev.map(server => 
                            server.uri === currentServer.uri 
                                ? { ...server, resources: resourcesData }
                                : server
                        ));
                        setCurrentServer(prev => ({ ...prev, resources: resourcesData }));
                    }
                    addMessage('system', `发现 ${resourcesData.length} 个资源`);
                }
                break;
            case 'call_tool':
                if (response.status === 'success') {
                    addMessage('tool', `工具执行结果: ${JSON.stringify(response.data, null, 2)}`);
                    message.success('工具调用成功');
                } else {
                    addMessage('error', `工具调用失败: ${response.message}`);
                    message.error('工具调用失败');
                }
                setLoading(false);
                break;
            case 'read_resource':
                if (response.status === 'success') {
                    addMessage('system', `资源内容:\n${JSON.stringify(response.data, null, 2)}`);
                    message.success('资源读取成功');
                } else {
                    addMessage('error', `资源读取失败: ${response.message}`);
                    message.error('资源读取失败');
                }
                setLoading(false);
                break;
            case 'error':
                addMessage('error', response.message);
                message.error('操作失败');
                setLoading(false);
                break;
            default:
                console.log('未知消息类型:', response);
        }
    };

    const handleDisconnect = async () => {
        setLoading(true);
        try {
            if (websocket && currentServer?.sessionId) {
                websocket.send(JSON.stringify({ 
                    type: 'disconnect',
                    sessionId: currentServer.sessionId
                }));
            }
        } catch (error) {
            addMessage('error', `断开连接错误: ${error.message}`);
            message.error('断开连接失败');
            setLoading(false);
        }
    };

    const handleDisconnectServer = async (server) => {
        setLoading(true);
        try {
            if (websocket && server.sessionId) {
                websocket.send(JSON.stringify({ 
                    type: 'disconnect',
                    sessionId: server.sessionId
                }));
                // 立即从列表中移除该服务器
                setConnectedServers(prev => prev.filter(s => s.sessionId !== server.sessionId));
                // 如果断开的是当前服务器，清空当前状态
                if (currentServer?.sessionId === server.sessionId) {
                    setCurrentServer(null);
                    setTools([]);
                    setResources([]);
                    setConnectionStatus(false);
                }
                addMessage('system', `已断开服务器: ${server.name}`);
                message.success('断开连接成功');
            }
        } catch (error) {
            addMessage('error', `断开连接错误: ${error.message}`);
            message.error('断开连接失败');
        }
        setLoading(false);
    };

    const handleSelectServer = (server) => {
        setCurrentServer(server);
        setTools(server.tools || []);
        setResources(server.resources || []);
        setConnectionStatus(true);
        addMessage('system', `已切换到服务器: ${server.name}`);
    };

    const loadTools = (sessionId = null) => {
        console.log('loadTools called with sessionId:', sessionId);
        console.log('currentServer:', currentServer);
        console.log('currentServer?.sessionId:', currentServer?.sessionId);
        
        const targetSessionId = sessionId || currentServer?.sessionId;
        console.log('targetSessionId:', targetSessionId);
        console.log('websocket readyState:', websocket?.readyState);
        console.log('WebSocket.OPEN:', WebSocket.OPEN);
        
        if (websocket && websocket.readyState === WebSocket.OPEN && targetSessionId) {
            console.log('Sending list_tools request with sessionId:', targetSessionId);
            websocket.send(JSON.stringify({ 
                type: 'list_tools',
                sessionId: targetSessionId
            }));
        } else {
            console.log('Cannot send list_tools request. Conditions check:');
            console.log('- websocket exists:', !!websocket);
            console.log('- websocket ready:', websocket?.readyState === WebSocket.OPEN);
            console.log('- targetSessionId exists:', !!targetSessionId);
        }
    };



    const loadResources = (sessionId = null) => {
        console.log('loadResources called with sessionId:', sessionId);
        console.log('currentServer:', currentServer);
        console.log('currentServer?.sessionId:', currentServer?.sessionId);
        
        const targetSessionId = sessionId || currentServer?.sessionId;
        console.log('targetSessionId:', targetSessionId);
        console.log('websocket readyState:', websocket?.readyState);
        console.log('WebSocket.OPEN:', WebSocket.OPEN);
        
        if (websocket && websocket.readyState === WebSocket.OPEN && targetSessionId) {
            console.log('Sending list_resources request with sessionId:', targetSessionId);
            websocket.send(JSON.stringify({ 
                type: 'list_resources',
                sessionId: targetSessionId
            }));
        } else {
            console.log('Cannot send list_resources request. Conditions check:');
            console.log('- websocket exists:', !!websocket);
            console.log('- websocket ready:', websocket?.readyState === WebSocket.OPEN);
            console.log('- targetSessionId exists:', !!targetSessionId);
        }
    };

    const handleCallTool = (toolName, params) => {
        setLoading(true);
        try {
            const parsedParams = JSON.parse(params);
            if (websocket && websocket.readyState === WebSocket.OPEN && currentServer?.sessionId) {
                websocket.send(JSON.stringify({
                    type: 'call_tool',
                    sessionId: currentServer.sessionId,
                    toolName: toolName,
                    arguments: parsedParams
                }));
            } else {
                addMessage('error', 'WebSocket连接未建立或sessionId缺失');
                message.error('连接错误');
                setLoading(false);
            }
        } catch (error) {
            addMessage('error', `工具调用错误: ${error.message}`);
            message.error('工具调用错误');
            setLoading(false);
        }
    };

    const handleUserInput = (command) => {
        if (!currentServer) {
            addMessage('error', '请先选择一个已连接的服务器');
            return;
        }
        
        addMessage('user', command);
        
        if (command === 'help') {
            addMessage('system', '可用的MCP命令:\n- list_tools: 列出可用工具\n- list_resources: 列出可用资源\n- call_tool: 调用指定工具\n- read_resource: 读取指定资源\n- help: 显示帮助信息');
            return;
        }
        
        if (websocket && websocket.readyState === WebSocket.OPEN) {
            if (command === 'list_tools') {
                websocket.send(JSON.stringify({
                    type: 'list_tools',
                    sessionId: currentServer.sessionId
                }));
            } else if (command === 'list_resources') {
                websocket.send(JSON.stringify({
                    type: 'list_resources',
                    sessionId: currentServer.sessionId
                }));
            } else if (command === 'call_tool') {
                // 显示工具选择和参数输入对话框
                Modal.confirm({
                    title: '调用工具',
                    width: 600,
                    content: (
                        <div>
                            <Form layout="vertical">
                                <Form.Item label="选择工具">
                                    <Select
                                        placeholder="选择要调用的工具"
                                        style={{ width: '100%' }}
                                        onChange={(value) => setSelectedTool(tools.find(t => t.name === value))}
                                        options={tools.map(tool => ({
                                            value: tool.name,
                                            label: `${tool.name} - ${tool.description || '无描述'}`
                                        }))}
                                    />
                                </Form.Item>
                                <Form.Item label="参数 (JSON格式)">
                                    <TextArea 
                                        rows={4}
                                        placeholder='输入JSON参数，例如: {"query": "test"}'
                                        defaultValue={toolParams}
                                        onChange={(e) => setToolParams(e.target.value)}
                                    />
                                </Form.Item>
                            </Form>
                        </div>
                    ),
                    onOk: () => {
                        if (!selectedTool) {
                            message.error('请选择要调用的工具');
                            return;
                        }
                        handleCallTool(selectedTool.name, toolParams);
                    }
                });
            } else if (command === 'read_resource') {
                // 显示资源选择对话框
                Modal.confirm({
                    title: '读取资源',
                    width: 600,
                    content: (
                        <div>
                            <Form layout="vertical">
                                <Form.Item label="选择资源">
                                    <Select
                                        placeholder="选择要读取的资源"
                                        style={{ width: '100%' }}
                                        onChange={(value) => setToolParams(value)}
                                        options={resources.map(resource => ({
                                            key: resource.uri,
                                            value: resource.uri,
                                            label: `${resource.name} - ${resource.uri}`
                                        }))}
                                    />
                                </Form.Item>
                                <Form.Item label="或手动输入资源URI">
                                    <Input 
                                        placeholder="输入资源URI"
                                        onChange={(e) => setToolParams(e.target.value)}
                                    />
                                </Form.Item>
                            </Form>
                        </div>
                    ),
                    onOk: () => {
                        if (!toolParams.trim()) {
                            message.error('请选择资源或输入资源URI');
                            return;
                        }
                        websocket.send(JSON.stringify({
                            type: 'read_resource',
                            sessionId: currentServer.sessionId,
                            resourceUri: toolParams
                        }));
                    }
                });
            }
        }
    };

    // 清理WebSocket连接
    useEffect(() => {
        return () => {
            if (websocket) {
                websocket.close();
            }
        };
    }, [websocket]);



    const sidebarItems = [
        {
            key: '1',
            label: (
                <div className="status-badge">
                    <div className={`status-dot ${connectionStatus ? 'connected' : ''}`}></div>
                    服务器管理
                </div>
            ),
            children: (
                <Space direction="vertical" style={{ width: '100%' }}>
                    <Form layout="vertical">
                        <Form.Item label="服务名">
                            <Input 
                                value={serverConfig.serverName}
                                onChange={(e) => setServerConfig({...serverConfig, serverName: e.target.value})}
                                placeholder="输入服务名称"
                            />
                        </Form.Item>
                        <Form.Item label="MCP服务器地址">
                            <Input 
                                value={serverConfig.serverUri}
                                onChange={(e) => setServerConfig({...serverConfig, serverUri: e.target.value})}
                                placeholder="ws://localhost:9000"
                            />
                        </Form.Item>
                        
                        <Form.Item>
                            <Alert 
                                message="说明" 
                                description="前端通过WebSocket连接到MCP客户端(ws://localhost:8080/ws/mcp)，然后由MCP客户端连接到指定的MCP服务器地址。" 
                                type="info" 
                                showIcon 
                                style={{ fontSize: '12px' }}
                            />
                        </Form.Item>
                    </Form>
                    
                    <Button 
                        type="primary" 
                        onClick={handleConnect}
                        loading={loading}
                        disabled={!serverConfig.serverName || !serverConfig.serverUri}
                        style={{ width: '100%' }}
                    >
                        连接新服务器
                    </Button>
                    
                    {connectedServers.length > 0 && (
                        <div>
                            <Divider style={{ margin: '12px 0' }}>已连接服务器</Divider>
                            {connectedServers.map((server, index) => (
                                <Card 
                                    key={server.sessionId} 
                                    size="small" 
                                    style={{ 
                                        marginBottom: '8px',
                                        cursor: 'pointer',
                                        border: currentServer?.sessionId === server.sessionId ? '2px solid #1890ff' : '1px solid #d9d9d9'
                                    }}
                                    onClick={() => handleSelectServer(server)}
                                >
                                    <div style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
                                        <div style={{ flex: 1 }}>
                                            <Text strong>{server.name}</Text>
                                            <br/>
                                            <Text type="secondary" style={{ fontSize: '11px' }}>
                                                {server.uri}
                                            </Text>
                                            <br/>
                                            <Text type="secondary" style={{ fontSize: '10px' }}>
                                                连接时间: {server.connectedAt}
                                            </Text>
                                        </div>
                                        <div style={{ display: 'flex', flexDirection: 'column', alignItems: 'center', gap: '4px' }}>
                                            <Badge 
                                                status={currentServer?.sessionId === server.sessionId ? 'processing' : 'default'} 
                                                text={currentServer?.sessionId === server.sessionId ? '当前' : '已连接'}
                                            />
                                            <Button 
                                                size="small" 
                                                danger
                                                onClick={(e) => {
                                                    e.stopPropagation();
                                                    handleDisconnectServer(server);
                                                }}
                                                loading={loading}
                                            >
                                                断开
                                            </Button>
                                        </div>
                                    </div>
                                </Card>
                            ))}
                        </div>
                    )}
                </Space>
            )
        },
        {
            key: '2',
            label: `工具列表 (${tools.length})`,
            children: (
                <Space direction="vertical" style={{ width: '100%' }}>
                    {tools.map(tool => {
                        // 解析inputSchema生成参数模板
                        const generateParamsTemplate = (inputSchema) => {
                            if (!inputSchema || !inputSchema.properties) return '{}';
                            
                            const params = {};
                            Object.keys(inputSchema.properties).forEach(key => {
                                const prop = inputSchema.properties[key];
                                switch (prop.type) {
                                    case 'string':
                                        params[key] = prop.default || '';
                                        break;
                                    case 'number':
                                    case 'integer':
                                        params[key] = prop.default || 0;
                                        break;
                                    case 'boolean':
                                        params[key] = prop.default || false;
                                        break;
                                    case 'array':
                                        params[key] = prop.default || [];
                                        break;
                                    case 'object':
                                        params[key] = prop.default || {};
                                        break;
                                    default:
                                        params[key] = prop.default || null;
                                }
                            });
                            return JSON.stringify(params, null, 2);
                        };
                        
                        const paramsTemplate = generateParamsTemplate(tool.inputSchema);
                        
                        return (
                            <Card key={tool.name} size="small" style={{ cursor: 'pointer' }}
                                  onClick={() => {
                                      setSelectedTool(tool);
                                      setToolParams(paramsTemplate);
                                      Modal.confirm({
                                          title: `调用工具: ${tool.name}`,
                                          width: 600,
                                          content: (
                                              <div>
                                                  <p><strong>描述:</strong> {tool.description || '无描述'}</p>
                                                  {tool.inputSchema && tool.inputSchema.properties && (
                                                      <div style={{ marginBottom: '16px' }}>
                                                          <p><strong>参数说明:</strong></p>
                                                          <ul style={{ fontSize: '12px', color: '#666' }}>
                                                              {Object.keys(tool.inputSchema.properties).map(key => {
                                                                  const prop = tool.inputSchema.properties[key];
                                                                  const required = tool.inputSchema.required?.includes(key) ? ' (必填)' : ' (可选)';
                                                                  return (
                                                                      <li key={key}>
                                                                          <strong>{key}</strong> ({prop.type}){required}: {prop.description || '无描述'}
                                                                      </li>
                                                                  );
                                                              })}
                                                          </ul>
                                                      </div>
                                                  )}
                                                  <p><strong>参数 (JSON格式):</strong></p>
                                                  <TextArea 
                                                      rows={6}
                                                      placeholder="输入JSON参数"
                                                      defaultValue={paramsTemplate}
                                                      onChange={(e) => setToolParams(e.target.value)}
                                                  />
                                              </div>
                                          ),
                                          onOk: () => handleCallTool(tool.name, toolParams)
                                      });
                                  }}>
                                <div>
                                    <Text strong>{tool.name}</Text>
                                    <br/>
                                    <Text type="secondary" style={{ fontSize: '12px' }}>
                                        {tool.description || '无描述'}
                                    </Text>
                                    {tool.inputSchema && tool.inputSchema.properties && (
                                        <div style={{ marginTop: '4px' }}>
                                            <Text type="secondary" style={{ fontSize: '10px' }}>
                                                参数: {Object.keys(tool.inputSchema.properties).join(', ')}
                                            </Text>
                                        </div>
                                    )}
                                </div>
                            </Card>
                        );
                    })}
                    {tools.length === 0 && (
                        <Text type="secondary">请先连接服务器</Text>
                    )}
                </Space>
            )
        },
        {
            key: '3',
            label: `资源列表 (${resources.length})`,
            children: (
                <Space direction="vertical" style={{ width: '100%' }}>
                    {resources.map(resource => (
                        <Card key={resource.uri} size="small">
                            <Text strong>{resource.name}</Text>
                            <br/>
                            <Text type="secondary" style={{ fontSize: '12px' }}>
                                {resource.uri}
                            </Text>
                        </Card>
                    ))}
                    {resources.length === 0 && (
                        <Text type="secondary">请先连接服务器</Text>
                    )}
                </Space>
            )
        }
    ];

    return (
        <div className="chatbox-container">
            <div className="chatbox-header">
                <Title level={2} style={{ color: 'white', margin: 0 }}>
                    MCP Client - 聊天界面
                </Title>
            </div>
            
            <div className="chatbox-main">
                <div className="chatbox-sidebar">
                    <Collapse 
                        items={sidebarItems}
                        defaultActiveKey={['1']}
                        style={{ border: 'none' }}
                    />
                </div>
                
                <div className="chatbox-chat-area">
                    <div className="chatbox-messages">
                        {messages.map(msg => (
                            <div key={msg.id} className={`message-item message-${msg.type}`}>
                                <div className="message-header">
                                    <Text strong>{msg.type.toUpperCase()}</Text>
                                    <Text type="secondary" style={{ fontSize: '12px', marginLeft: '8px' }}>
                                        {msg.timestamp}
                                    </Text>
                                </div>
                                <div className="message-content">
                                    <pre style={{ whiteSpace: 'pre-wrap', margin: 0 }}>{msg.content}</pre>
                                </div>
                            </div>
                        ))}
                        <div ref={messagesEndRef} />
                        {messages.length === 0 && (
                            <div style={{ textAlign: 'center', padding: '40px', color: '#999' }}>
                                <Text type="secondary">欢迎使用MCP客户端！请先在左侧配置并连接服务器，然后选择服务器并输入 MCP 命令进行交互。</Text>
                            </div>
                        )}
                    </div>
                    
                    <div className="chatbox-input">
                        <Space.Compact style={{ width: '100%' }}>
                            <Select
                                value={inputMessage}
                                onChange={(value) => setInputMessage(value)}
                                placeholder="选择MCP命令..."
                                style={{ width: '100%' }}
                                disabled={connectedServers.length === 0 || !currentServer}
                                options={[
                                    { value: 'list_tools', label: 'list_tools - 列出可用工具' },
                                    { value: 'list_resources', label: 'list_resources - 列出可用资源' },
                                    { value: 'call_tool', label: 'call_tool - 调用指定工具' },
                                    { value: 'read_resource', label: 'read_resource - 读取指定资源' },
                                    { value: 'help', label: 'help - 显示帮助信息' }
                                ]}
                            />
                            <Button 
                                type="primary" 
                                onClick={() => {
                                    if (inputMessage) {
                                        handleUserInput(inputMessage);
                                        setInputMessage('');
                                    }
                                }}
                                disabled={connectedServers.length === 0 || !currentServer || !inputMessage}
                            >
                                执行
                            </Button>
                        </Space.Compact>
                    </div>
                </div>
            </div>
        </div>
    );
};

export default ChatBox;